import os
import datetime
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from sklearn.preprocessing import StandardScaler
import joblib
from util.logUtil import Logger
from util.commonUtil import data_preprocessing


def analysis_data(data, logger):
    """
    数据分析：计算数值特征与目标变量（Attrition）的AUC值，评估特征相关性
    :param data: 输入DataFrame（包含'Attrition'列）
    :param logger: 日志实例
    :return: 数值特征的AUC值列表
    """
    logger.info("=========开始数据分析==========")
    numeric_cols = data.select_dtypes(include=['int64', 'float64']).columns.tolist()
    numeric_cols = [col for col in numeric_cols if col != 'Attrition']  # 排除目标列
    logger.info(f"待分析的数值特征：{numeric_cols}")

    auc_list = []
    for col in numeric_cols:
        try:
            # 计算特征与目标变量的AUC（评估区分能力）
            from sklearn.metrics import roc_auc_score
            auc_score = roc_auc_score(data['Attrition'], data[col])
            auc_list.append(auc_score)
            logger.info(f"特征'{col}'的AUC值：{auc_score:.4f}")
        except Exception as e:
            logger.warning(f"计算特征'{col}'的AUC失败：{str(e)}")

    logger.info("=========数据分析结束==========")
    return auc_list


def data_processing(data, logger, is_train=True, scaler=None):
    """
    特征工程：处理分类变量、创建新特征、标准化特征
    :param data: 输入DataFrame（包含'Attrition'列）
    :param logger: 日志实例
    :param is_train: 是否为训练数据（True时拟合标准化器，False时使用已有标准化器）
    :param scaler: 标准化器（训练时为None，预测时传入训练好的scaler）
    :return: x（处理后的特征）, y（目标变量）, scaler（训练时返回拟合好的scaler）
    """
    logger.info("=========开始特征工程==========")

    # 1. 分类变量映射（将字符串转换为数值）
    category_mapping = {
        'BusinessTravel': {'Non-Travel': 0, 'Travel_Rarely': 1, 'Travel_Frequently': 2},
        'Department': {'Human Resources': 1, 'Research & Development': 2, 'Sales': 3},
        'MaritalStatus': {'Divorced': 0, 'Single': 1, 'Married': 2},
        'Gender': {'Male': 0, 'Female': 1}
    }

    for col, mapping in category_mapping.items():
        if col in data.columns:
            data[col] = data[col].map(mapping)
            logger.info(f"分类特征'{col}'映射完成，映射规则：{mapping}")
        else:
            logger.warning(f"特征'{col}'不存在于数据中，跳过映射")

    # 2. 独热编码（处理其他分类变量）
    data_dum = pd.get_dummies(data, drop_first=True).astype(int)  # drop_first避免多重共线性
    logger.info(f"独热编码后数据形状：{data_dum.shape}")

    # 3. 特征组合（创建高价值衍生特征）
    # 组合满意度特征
    if all(col in data_dum.columns for col in
           ['EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction']):
        data_dum['STA'] = data_dum['EnvironmentSatisfaction'] + data_dum['JobSatisfaction'] + data_dum[
            'RelationshipSatisfaction']
        logger.info("创建组合特征'STA'（满意度总和）")

    # 工作满意度/工作投入（加1避免除零）
    if all(col in data_dum.columns for col in ['JobSatisfaction', 'JobInvolvement']):
        data_dum['job_ment_sta'] = (data_dum['JobSatisfaction'] + 1) / (data_dum['JobInvolvement'] + 1)
        logger.info("创建组合特征'job_ment_sta'（工作满意度/工作投入）")

    # 平均每家公司工作年限
    if all(col in data_dum.columns for col in ['NumCompaniesWorked', 'TotalWorkingYears']):
        data_dum['avg_year_work'] = (data_dum['NumCompaniesWorked'] + 1) / (data_dum['TotalWorkingYears'] + 1)
        logger.info("创建组合特征'avg_year_work'（平均每家公司工作年限）")

    # 4. 删除冗余特征（无区分度或已被组合的特征）
    drop_cols = [
        'EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction',
        'NumCompaniesWorked', 'TotalWorkingYears', 'Over18_Y', 'StandardHours', 'EmployeeNumber'
    ]
    drop_cols = [col for col in drop_cols if col in data_dum.columns]  # 只删除存在的列
    data_dum.drop(drop_cols, axis=1, inplace=True)
    logger.info(f"删除冗余特征：{drop_cols}，处理后数据形状：{data_dum.shape}")

    # 5. 提取特征和目标变量
    x = data_dum.drop('Attrition', axis=1)
    y = data_dum['Attrition']
    logger.info(f"特征矩阵形状：{x.shape}，目标变量形状：{y.shape}")

    # 6. 特征标准化（训练时拟合，预测时复用）
    if is_train:
        scaler = StandardScaler()
        x = scaler.fit_transform(x)
        logger.info("训练数据标准化完成（已拟合StandardScaler）")
    else:
        if scaler is None:
            raise ValueError("预测时必须传入训练好的StandardScaler")
        x = scaler.transform(x)
        logger.info("测试数据标准化完成（使用训练集StandardScaler）")

    logger.info("=========特征工程结束==========")
    return x, y, scaler


def model_train(x_train, y_train, scaler, logger):
    """
    模型训练与调优：使用网格搜索优化逻辑回归模型
    :param x_train: 训练特征（已标准化）
    :param y_train: 训练目标变量
    :param scaler: 训练好的标准化器（需保存供预测使用）
    :param logger: 日志实例
    :return: 最佳模型
    """
    logger.info("=========开始模型训练与调优==========")

    # 定义超参数网格（逻辑回归关键参数）
    param_grid = {
        'C': [0.01, 0.1, 1, 10],  # 正则化强度（倒数），越小正则化越强
        'penalty': ['l1', 'l2'],  # 正则化方式（l1:Lasso, l2:Ridge）
        'solver': ['liblinear'],  # 仅liblinear支持l1正则化
        'class_weight': ['balanced', None],  # 类别权重（处理不平衡数据）
        'max_iter': [1000, 2000]  # 迭代次数
    }
    logger.info(f"超参数搜索网格：{param_grid}")

    # 初始化模型与网格搜索
    base_model = LogisticRegression(random_state=42)  # 固定随机种子确保可复现
    grid_search = GridSearchCV(
        estimator=base_model,
        param_grid=param_grid,
        cv=5,  # 5折交叉验证
        scoring='roc_auc',  # 以AUC为优化目标
        n_jobs=-1,  # 并行计算（使用所有CPU核心）
        verbose=1
    )

    # 执行网格搜索
    grid_search.fit(x_train, y_train)
    logger.info(f"最佳交叉验证AUC：{grid_search.best_score_:.4f}")
    logger.info(f"最佳超参数：{grid_search.best_params_}")

    # 保存最佳模型和标准化器
    model_dir = '../model'
    os.makedirs(model_dir, exist_ok=True)
    joblib.dump(grid_search.best_estimator_, os.path.join(model_dir, 'best_model.pkl'))
    joblib.dump(scaler, os.path.join(model_dir, 'scaler.pkl'))
    logger.info(f"最佳模型保存至：{os.path.join(model_dir, 'best_model.pkl')}")
    logger.info(f"标准化器保存至：{os.path.join(model_dir, 'scaler.pkl')}")

    logger.info("=========模型训练与调优结束==========")
    return grid_search.best_estimator_


class PowerLoadModel(object):
    """此处实际为员工流失预测，保持类名兼容性"""

    def __init__(self, filename):
        """
        初始化模型：创建日志、加载并预处理数据
        :param filename: 数据文件路径
        """
        # 1. 初始化日志（区分训练/预测日志）
        log_type = 'train' if 'train' in filename.lower() else 'predict'
        log_name = f"{log_type}_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}"
        self.logfile = Logger('../', log_name).get_logger()
        self.logfile.info(f"初始化PowerLoadModel，数据文件：{filename}")

        # 2. 加载并预处理数据
        self.data_source = data_preprocessing(filename, self.logfile)
        self.logfile.info(f"数据加载完成，样本量：{len(self.data_source)}，特征数：{len(self.data_source.columns)}")